import os
import time
import torch
import numpy as np
from torch.autograd import Variable
import torch.nn as nn
from utils.pointnet2_utils import furthest_point_sample, gather_operation
from common.lib.metric import to_python_float, reduce_tensor, AverageMeter, accuracy
import time


def validate(val_loader, model, epoch, config, logger):
    stage_one_grasp_cls_loss = AverageMeter()
    stage_one_grasp_score_loss = AverageMeter()
    stage_one_suction_cls_loss = AverageMeter()
    stage_one_suction_score_loss = AverageMeter()
    stage_one_objectness_loss = AverageMeter()
    stage_one_grasp_pose_loss = AverageMeter()
    stage_one_normal_loss = AverageMeter()
    stage_one_distance_loss = AverageMeter()
    stage_one_suction_pose_loss = AverageMeter()
    stage_two_cls_loss = AverageMeter()
    top1 = AverageMeter()
    # top2 = AverageMeter()
    iou = AverageMeter()
    recall = AverageMeter()

    iou_suction = AverageMeter()
    recall_suction = AverageMeter()

    n_iteration_per_epoch = len(val_loader)
    model.eval()
    # model.train()

    with torch.no_grad():
        end = time.time()
        for i, batch in enumerate(val_loader):
            data_time = time.time() - end
            net_time = time.time()
            data, label, dataset_idx = batch
            data = data.cuda()
            if config.train_grasp:
                if config.train_grasp_approach:
                    grasp_cls_label = label[0].cuda(non_blocking=True)
                    grasp_score_label = label[1].cuda(non_blocking=True)
                    grasp_pose_label = label[2].cuda(non_blocking=True)
                    matrix_label = label[3].cuda(non_blocking=True)
                    distance_label = label[4].cuda(non_blocking=True)
                    normals_label = label[5].cuda(non_blocking=True)
                    suction_cls_label = label[6].cuda(non_blocking=True)
                    suction_score_label = label[7].cuda(non_blocking=True)
                    objectness_label = label[8].cuda(non_blocking=True)

                else:
                    grasp_cls_label = label[0].cuda(non_blocking=True)
                    grasp_score_label = label[1].cuda(non_blocking=True)
                    grasp_pose_label = label[2].cuda(non_blocking=True)
                    matrix_label = label[3].cuda(non_blocking=True)

            elif config.train_suction_only:
                suction_cls_label = label[0].cuda(non_blocking=True)
                suction_score_label = label[1].cuda(non_blocking=True)
                normals_label = label[2].cuda(non_blocking=True)
                objectness_label = label[3].cuda(non_blocking=True)


            # if not config.non_uniform_sampling:
            #     idx = furthest_point_sample(data, int(config.num_points / config.sample_rate))
            #     data = gather_operation(
            #         data.transpose(1, 2).contiguous(),
            #         idx).transpose(1, 2).contiguous()
            #     cls_label = gather_operation(
            #         cls_label.float().unsqueeze(dim=2).transpose(1, 2).contiguous(),
            #         idx).squeeze().long().contiguous()
            #     pose_label = gather_operation(
            #         pose_label.transpose(1, 2).contiguous(),
            #         idx).transpose(1, 2).contiguous()

            if config.train_grasp:

                if config.train_grasp_approach and config.train_suction_combine:

                    kwargs = {'grasp_cls_label': grasp_cls_label, 'grasp_score_label': grasp_score_label,
                              'grasp_pose_label': grasp_pose_label,
                              'matrix_label': matrix_label, 'distance_label': distance_label,
                              'normals_label': normals_label,
                              'suction_cls_label': suction_cls_label, 'suction_score_label': suction_score_label,
                              'objectness_label': objectness_label, 'mode': 'train'}

                else:
                    kwargs = {'grasp_cls_label': grasp_cls_label, 'grasp_score_label': grasp_score_label,
                              'grasp_pose_label': grasp_pose_label, 'matrix_label': matrix_label, 'mode': 'train'}

            elif config.train_suction_only:
                kwargs = {'suction_cls_label': suction_cls_label, 'score_label': suction_score_label,
                          'normals_label': normals_label,
                          'objectness_label': objectness_label, 'mode': 'train'}

            batch_dict = model(data, **kwargs)
            # batch_dict_, cls_pred, _, _ = model.module.back_bone(data, cls_label, pose_label)
            # print(cls_pred.shape)
            # exit()
            net_time = time.time() - net_time

            metric_time = time.time()
            if not config.use_focalloss:
                stage_one_prec1 = accuracy(batch_dict['batch_cls_preds'].data, grasp_cls_label, topk=(1,))
                stage_one_prec1 = stage_one_prec1[0]
            if config.distributed:
                stage_one_reduced_grasp_cls_loss = reduce_tensor(batch_dict['stage_one_grasp_cls_loss'].data)
                stage_one_reduced_grasp_score_loss = reduce_tensor(batch_dict['stage_one_grasp_score_loss'].data)
                stage_one_reduced_suction_cls_loss = reduce_tensor(batch_dict['stage_one_suction_cls_loss'].data)
                stage_one_reduced_suction_score_loss = reduce_tensor(batch_dict['stage_one_suction_score_loss'].data)
                stage_one_reduced_objectness_loss = reduce_tensor(batch_dict['stage_one_objectness_loss'].data)

                if config.train_grasp:
                    if config.train_grasp_approach and config.train_suction_combine:
                        stage_one_reduced_grasp_pose_loss = reduce_tensor(batch_dict['stage_one_grasp_pose_loss'].data)
                        stage_one_reduced_normal_loss = reduce_tensor(batch_dict['stage_one_normal_loss'].data)
                        stage_one_reduced_distance_loss = reduce_tensor(batch_dict['stage_one_distance_loss'].data)
                    else:
                        stage_one_reduced_grasp_pose_loss = reduce_tensor(batch_dict['stage_one_grasp_pose_loss'].data)
                if config.train_suction_only:
                    stage_one_reduced_suction_pose_loss = reduce_tensor(batch_dict['stage_one_suction_pose_loss'].data)
                if config.train_two_stage:
                    stage_two_reduced_cls_loss = reduce_tensor(batch_dict['stage_two_loss'].data)

                if not config.use_focalloss:
                    stage_one_prec1 = reduce_tensor(stage_one_prec1)
                    # prec2 = reduce_tensor(prec2)
            else:
                stage_one_reduced_grasp_cls_loss = reduce_tensor(batch_dict['stage_one_grasp_cls_loss'].data)
                stage_one_reduced_grasp_score_loss = reduce_tensor(batch_dict['stage_one_grasp_score_loss'].data)
                stage_one_reduced_suction_cls_loss = reduce_tensor(batch_dict['stage_one_suction_cls_loss'].data)
                stage_one_reduced_suction_score_loss = reduce_tensor(batch_dict['stage_one_suction_score_loss'].data)
                stage_one_reduced_objectness_loss = reduce_tensor(batch_dict['stage_one_objectness_loss'].data)

                if config.train_grasp:
                    if config.train_grasp_approach and config.train_suction_combine:
                        stage_one_reduced_grasp_pose_loss = reduce_tensor(batch_dict['stage_one_grasp_pose_loss'].data)
                        stage_one_reduced_normal_loss = reduce_tensor(batch_dict['stage_one_normal_loss'].data)
                        stage_one_reduced_distance_loss = reduce_tensor(batch_dict['stage_one_distance_loss'].data)
                    else:
                        stage_one_reduced_grasp_pose_loss = reduce_tensor(batch_dict['stage_one_grasp_pose_loss'].data)
                if config.train_suction_only:
                    stage_one_reduced_suction_pose_loss = reduce_tensor(batch_dict['stage_one_suction_pose_loss'].data)
                if config.train_two_stage:
                    stage_two_reduced_cls_loss = reduce_tensor(batch_dict['stage_two_loss'].data)

            stage_one_grasp_cls_loss.update(to_python_float(stage_one_reduced_grasp_cls_loss), data.size(0))
            stage_one_grasp_score_loss.update(to_python_float(stage_one_reduced_grasp_score_loss), data.size(0))
            stage_one_suction_cls_loss.update(to_python_float(stage_one_reduced_suction_cls_loss), data.size(0))
            stage_one_suction_score_loss.update(to_python_float(stage_one_reduced_suction_score_loss), data.size(0))
            stage_one_objectness_loss.update(to_python_float(stage_one_reduced_objectness_loss), data.size(0))

            if config.train_grasp:
                if config.train_grasp_approach and config.train_suction_combine:
                    stage_one_grasp_pose_loss.update(to_python_float(stage_one_reduced_grasp_pose_loss), data.size(0))
                    stage_one_normal_loss.update(to_python_float(stage_one_reduced_normal_loss), data.size(0))
                    stage_one_distance_loss.update(to_python_float(stage_one_reduced_distance_loss), data.size(0))
                else:
                    stage_one_grasp_pose_loss.update(to_python_float(stage_one_reduced_grasp_pose_loss), data.size(0))
            if config.train_suction_only:
                stage_one_suction_pose_loss.update(to_python_float(stage_one_reduced_suction_pose_loss), data.size(0))
            if config.train_two_stage:
                stage_two_cls_loss.update(to_python_float(stage_two_reduced_cls_loss), data.size(0))
            if not config.use_focalloss:
                top1.update(to_python_float(stage_one_prec1), data.size(0))
                # top2.update(to_python_float(prec2), data.size(0))
            metric_time = time.time() - metric_time


            if config.use_focalloss:
                grasp_cls_scores = batch_dict['batch_grasp_cls_preds'].squeeze()
                suction_cls_scores = batch_dict['batch_suction_cls_preds'].squeeze()
                # cls_pred = cls_pred.squeeze()
                # cls_scores = torch.sigmoid(cls_pred)
                for pred_, cls_ in zip(grasp_cls_scores, grasp_cls_label):
                    area_1 = np.where(pred_.detach().cpu() > config.seg_thresh)
                    area_2 = np.where(cls_.cpu() == 1)
                    area_1 = set(area_1[0])
                    area_2 = set(area_2[0])
                    U = set.union(area_1, area_2)
                    I = set.intersection(area_1, area_2)
                    if len(U) == 0:
                        iou_ = 0
                    else:
                        iou_ = len(I) / len(U)
                    if len(area_2) == 0:
                        recall_ = 0
                    else:
                        recall_ = len(I) / len(area_2)
                    iou.update(iou_, 1)
                    recall.update(recall_, 1)

                for pred_s, cls_s in zip(suction_cls_scores, suction_cls_label):
                    area_1_s = np.where(pred_s.detach().cpu() > config.seg_thresh)
                    area_2_s = np.where(cls_s.cpu() == 1)
                    area_1_s = set(area_1_s[0])
                    area_2_s = set(area_2_s[0])
                    U_s = set.union(area_1_s, area_2_s)
                    I_s = set.intersection(area_1_s, area_2_s)
                    if len(U_s) == 0:
                        iou_s = 0
                    else:
                        iou_s = len(I_s) / len(U_s)
                    if len(area_2_s) == 0:
                        recall_s = 0
                    else:
                        recall_s = len(I_s) / len(area_2_s)
                    iou_suction.update(iou_s, 1)
                    recall_suction.update(recall_s, 1)


            if i % int(config.print_freq) == 0 and config.rank == 0:
                if not config.use_focalloss:
                    logger.info('Train:[{0}] [{1}/{2}]\t'
                                'Data_time{3:.3f}\t'
                                'Net_time {4:.3f}\t'
                                'Metric_time {5:.3f}\t'
                                'S-1 GraspClsloss: {stage1_grasp_cls_loss.avg:.4f}\t'
                                'S-1 GraspScoreloss: {stage1_grasp_score_loss.avg:.4f}\t'
                                'S-1 Objectness loss: {stage1_objectness_loss.avg:.4f}\t'
                                'Top1 is {top1.val:.4f} ({top1.avg:.4f})'.format(
                        epoch, i + 1, n_iteration_per_epoch,
                        data_time, net_time, metric_time,
                        stage1_grasp_cls_loss = stage_one_grasp_cls_loss,
                        stage1_grasp_score_loss = stage_one_grasp_score_loss,
                        stage1_objectness_loss = stage_one_objectness_loss,
                        top1 = top1))
                    if config.train_grasp:
                        if config.train_grasp_approach and config.train_suction_combine:
                            logger.info('S-1 Suction_cls_loss: {stage1_suction_cls_loss.avg:.6f}'.format(
                                stage1_suction_cls_loss=stage_one_suction_cls_loss))
                            logger.info('S-1 Suction_score_loss: {stage1_suction_score_loss.avg:.6f}'.format(
                                stage1_suction_score_loss=stage_one_suction_score_loss))
                            logger.info('S-1 Grasp_Pose_loss: {stage1_grasp_pose_loss.avg:.6f}'.format(
                                stage1_grasp_pose_loss=stage_one_grasp_pose_loss))
                            logger.info('S-1 Normal loss: {stage1_normal_loss.avg:.6f}'.format(
                                stage1_normal_loss=stage_one_normal_loss))
                            logger.info('S-1 Distance loss: {stage1_distance_loss.avg:.6f}'.format(
                                stage1_distance_loss=stage_one_distance_loss))
                        else:
                            logger.info('S-1 Poseloss: {stage1_grasp_pose_loss.avg:.6f}'.format(
                                stage1_grasp_pose_loss=stage_one_grasp_pose_loss))
                    if config.train_suction_only:
                        logger.info('S-1 Poseloss: {stage1_suction_pose_loss.avg:.6f}'.format(
                            stage1_suction_pose_loss=stage_one_suction_pose_loss))
                    if config.train_two_stage:
                        logger.info('S-2 Clsloss: {stage2_cls_loss.avg:.6f}'.format(stage2_cls_loss=stage_two_cls_loss))
                else:
                    logger.info('Test:[{0}] [{1}/{2}]\t'
                                'Data_time{3:.3f}\t'
                                'Net_time {4:.3f}\t'
                                'Metric_time {5:.3f}\t'
                                'Grasp_Iou {iou.val:.4f} ({iou.avg:.4f})\t'
                                'Grasp_Recall {recall.val:.4f} ({recall.avg:.4f})\t'
                                'Suction_Iou {iou_s.val:.4f} ({iou_s.avg:.4f})\t'
                                'Suction_Recall {recall_s.val:.4f} ({recall_s.avg:.4f})\t'
                                'S-1 Grasp_Cls_loss: {stage1_grasp_cls_loss.avg:.4f}\t'
                                'S-1 Grasp_Scoreloss:{stage1_grasp_score_loss.avg:.4f}\t'
                                'S-1 Objectness loss: {stage1_objectness_loss.avg:.4f}'
                        .format(
                        epoch, i + 1, n_iteration_per_epoch,
                        data_time, net_time, metric_time,
                        iou=iou, recall=recall,
                        iou_s = iou_suction, recall_s = recall_suction,
                        stage1_grasp_cls_loss = stage_one_grasp_cls_loss,
                        stage1_grasp_score_loss = stage_one_grasp_score_loss,
                        stage1_objectness_loss = stage_one_objectness_loss))
                    if config.train_grasp:
                        if config.train_grasp_approach and config.train_suction_combine:
                            logger.info('S-1 Suction_cls_loss: {stage1_suction_cls_loss.avg:.6f}'.format(
                                stage1_suction_cls_loss=stage_one_suction_cls_loss))
                            logger.info('S-1 Suction_score_loss: {stage1_suction_score_loss.avg:.6f}'.format(
                                stage1_suction_score_loss=stage_one_suction_score_loss))
                            logger.info('S-1 Grasp_Pose_loss: {stage1_grasp_pose_loss.avg:.6f}'.format(
                                stage1_grasp_pose_loss=stage_one_grasp_pose_loss))
                            logger.info('S-1 Normal loss: {stage1_normal_loss.avg:.6f}'.format(
                                stage1_normal_loss=stage_one_normal_loss))
                            logger.info('S-1 Distance loss: {stage1_distance_loss.avg:.6f}'.format(
                                stage1_distance_loss=stage_one_distance_loss))
                        else:
                            logger.info('S-1 Poseloss: {stage1_grasp_pose_loss.avg:.6f}'.format(
                                stage1_grasp_pose_loss=stage_one_grasp_pose_loss))
                    if config.train_suction_only:
                        logger.info('S-1 Poseloss: {stage1_suction_pose_loss.avg:.6f}'.format(
                            stage1_suction_pose_loss=stage_one_suction_pose_loss))
                    if config.train_two_stage:
                        logger.info(
                            'S-2 CLS loss: {stage2_cls_loss.avg:.6f}\t'.format(stage2_cls_loss=stage_two_cls_loss))

            end = time.time()

        if not config.use_focalloss:
            return top1.avg
        else:
            return stage_one_grasp_cls_loss.avg, stage_one_grasp_score_loss.avg, \
                   stage_one_suction_cls_loss.avg, stage_one_suction_score_loss.avg, \
                   stage_one_objectness_loss.avg, stage_one_grasp_pose_loss.avg, \
                   stage_one_normal_loss.avg, stage_one_distance_loss.avg

